import torch
import cv2
import numpy as np
from PIL import Image
import torch.nn.functional as F
from facenet_pytorch import MTCNN, InceptionResnetV1
import os
import time

# 确保设备可用
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')

# 定义 MTCNN 模型用于检测面部并裁剪
mtcnn = MTCNN(device=device)

# 定义预训练的 FaceNet 模型（InceptionResnetV1）
resnet = InceptionResnetV1(pretrained='vggface2').eval().to(device)

# 加载两张图像，构建集合
img_paths_collection = []
for i in os.listdir("baseface"):
    img_paths_collection.append("baseface/" + i)

imgs_collection = [Image.open(img_path) for img_path in img_paths_collection]

# 使用 MTCNN 检测面部并裁剪
cropped_imgs_collection = []
for img in imgs_collection:
    try:
        cropped_img = mtcnn(img)
        cropped_imgs_collection.append(cropped_img)
    except Exception as e:
        print(f"Error during face detection and cropping: {e}")

# 将裁剪后的图像放入一个 batch 中
batch_collection = torch.stack(cropped_imgs_collection).to(device)

# 使用 torch.no_grad() 进行特征提取，以节省内存和加速计算
with torch.no_grad():
    embeddings_collection = resnet(batch_collection)

# 加载第三张图片
third_img_path = "testface/degui2.jpg"
third_img = Image.open(third_img_path)

# 使用 MTCNN 检测面部并裁剪
try:
    cropped_third_img = mtcnn(third_img)
except Exception as e:
    print(f"Error during face detection and cropping: {e}")
else:
    # 将裁剪后的图像从张量转换为 NumPy 数组
    cropped_third_img_np = cropped_third_img.permute(1, 2, 0).cpu().numpy()

    # 将 NumPy 数组的值从 [0, 1] 转换为 [0, 255]
    cropped_third_img_np = np.clip(cropped_third_img_np * 255, 0, 255).astype(np.uint8)

    # 将 RGB 转换为 BGR
    cropped_third_img_np_bgr = cv2.cvtColor(cropped_third_img_np, cv2.COLOR_RGB2BGR)

    # 使用 cv2.imshow 显示图像
    cv2.imshow('Cropped Face', cropped_third_img_np_bgr)
    cv2.waitKey(0)
    cv2.destroyAllWindows()

    # 提取第三张图片的特征向量
    third_embedding = resnet(cropped_third_img.unsqueeze(0).to(device))

    # 计算第三张图片与集合中所有元素的余弦相似度
    similarities = F.cosine_similarity(third_embedding, embeddings_collection)

    # 找出相似度最高的元素
    start_time = time.time()
    max_similarity, max_index = similarities.max(dim=0)
    end_time = time.time()
    print(f"找出相似度最高图片花费时间为:{end_time - start_time}ms")
    # 输出相似度最高的元素
    print(f"The third image is most similar to the {max_index.item() + 1}th image in the collection.")
    print(f"img :{img_paths_collection[max_index]}")
    print(f"The highest cosine similarity is: {max_similarity.item()}")

    # 输出所有相似度（可选）
    # print("Similarities with all images in the collection:")
    # print(similarities.squeeze().cpu().numpy())
